import os
import re
import time
from typing import Union, Callable, List, Set, Dict

from scripts.Log import Log
from scripts.device.AdbShell import AdbShell, PmListPkgInfo
from scripts.util.Util import Util
from scripts.actions.Action import Action
from scripts.actions.Context import Context


class BackupAppAction(Action):
    KEY = Action.register('app')
    DISPLAY_NAME = 'Installed apps'

    PATTERN_APK = re.compile(r'(.+)\.[aA][pP][kK]\s*')

    def __init__(self, context: Context):
        super().__init__(BackupAppAction.KEY, BackupAppAction.DISPLAY_NAME, context)

    def _on_backup(self, progress: Union[None, Callable[[int, int, str], None]]):
        device = self.wait_for_boot_up_ready()
        with device.shell() as shell:
            pkg_infos = Util.dump_pkg_infos(shell)

            # start users
            user_list = pkg_infos.keys()
            Action.start_users(user_list, shell=shell)
            Action.wait_for_filesystem_ready(user_list, shell=shell)

            # make backup list
            backup_dict: Dict[str, BackupAppAction.BackupItem] = dict()
            for user in pkg_infos:
                user_pkg_infos = pkg_infos[user]
                for pkg in user_pkg_infos:
                    info = user_pkg_infos[pkg]
                    if not info.apk_path.startswith('/data/app/'):
                        continue  # ignore system apps

                    backup_item: BackupAppAction.BackupItem
                    if pkg in backup_dict:
                        backup_item = backup_dict[pkg]
                    else:
                        device_dir = os.path.dirname(info.apk_path)
                        paths = shell.find(device_dir, 1, 1)
                        paths = list(filter(lambda p: BackupAppAction.PATTERN_APK.fullmatch(p) is not None, paths))
                        backup_dict[pkg] = backup_item = BackupAppAction.BackupItem(pkg, paths, self)
                    backup_item.users.add(user)

            # do backup

            i = 0
            for pkg in backup_dict:
                if progress:
                    progress(i, len(backup_dict), pkg)
                i += 1

                # backup files
                item = backup_dict[pkg]
                for file_item in item.file_list:
                    Action.pull_file(device, file_item, shell)

                # backup installed users
                users = ['%d' % i for i in item.users]
                sorted(users)
                str_users = ' '.join(users)
                with open(item.host_user_path, 'w', encoding='utf8') as file:
                    file.write(str_users)

                Log.d('Backup: [%s: %s]' % (str_users, pkg))

            if progress:
                progress(len(backup_dict), len(backup_dict), 'Done [%s]' % self.duration_time_str())

            Log.i('Backup success for %d apps' % len(backup_dict))

    def _on_restore(self, progress: Union[None, Callable[[int, int, str], None]]):
        device = self.wait_for_boot_up_ready()
        with device.shell() as shell:
            # start users
            user_list = {user.user_id for user in shell.pm_list_users()}
            Action.start_users(user_list, shell=shell)
            Action.wait_for_filesystem_ready(user_list, shell=shell)

            # disable package verifier
            Log.d('Disable package verifier')
            shell.settings_put('global', 'package_verifier_enable', '0')

            # get restore list
            restore_dict: Dict[str, BackupAppAction.RestoreItem] = dict()
            for pkg in os.listdir(self.data_dir()):
                restore_dict[pkg] = item = BackupAppAction.RestoreItem(pkg)
                host_dir = os.path.join(self.data_dir(), pkg)

                # load users
                users: List[int]
                with open(os.path.join(host_dir, 'user.txt'), 'r', encoding='utf8') as file:
                    str_users = file.read()
                    users = [int(user) for user in str_users.split(' ')]
                    item.users = set(users) & user_list

                # load files
                for file_name in os.listdir(host_dir):
                    if BackupAppAction.PATTERN_APK.fullmatch(file_name):
                        item.host_apk_list.append(os.path.join(host_dir, file_name))

            # do restore
            i = 0
            for pkg in restore_dict:
                item = restore_dict[pkg]
                if progress:
                    progress(i, len(restore_dict), item.pkg)
                i += 1

                # install apk
                first_install = True
                for user in item.users:
                    if first_install:
                        first_install = False

                        # install apk
                        retry = 3
                        while True:
                            retry -= 1
                            session: AdbShell.InstallSession = shell.pm_install_session(user)
                            with session:
                                for host_apk_path in item.host_apk_list:
                                    device_apk_path = '%s/%s' % (self.device_temp_dir(),
                                                                 os.path.basename(host_apk_path))
                                    device.push(host_apk_path, device_apk_path)
                                    session.write(device_path=device_apk_path)
                                    shell.rm(device_apk_path, recursive=False)
                                    # session.write(host_path=host_apk_path)
                            if session.result or retry > 0:
                                break
                            else:
                                Log.w('install pkg [%s] failed, retry after 1 second' % pkg)
                                time.sleep(1)
                    else:
                        shell.pm_install_existing(user, pkg)

                Log.d('Restore: [%s: %s]' % (' '.join([str(user) for user in item.users]), item.pkg))

            if progress:
                progress(len(restore_dict), len(restore_dict), 'Done [%s]' % self.duration_time_str())

            Log.i('Restore success for %d apps' % len(restore_dict))

    @staticmethod
    def _find_pkg_info(infos: Dict[int, Dict[str, PmListPkgInfo]],
                       pkg: str) -> Union[None, PmListPkgInfo]:
        for user in infos:
            if pkg in infos[user]:
                return infos[user][pkg]
        return None

    class BackupItem(object):
        def __init__(self, pkg: str, device_file_list: List[str], action):
            action: BackupAppAction
            self.pkg = pkg
            self.device_file_list = device_file_list
            self.users: Set[int] = set()
            self.host_user_path = os.path.join(action.data_dir(), self.pkg, 'user.txt')
            self.file_list = [
                Action.FileTransferItem(
                    device_file,
                    os.path.join(action.data_dir(), self.pkg, os.path.basename(device_file)),
                    os.path.join(action.device_temp_dir(), os.path.basename(device_file)),
                    compress=False, clear_copy=False
                )
                for device_file in self.device_file_list
            ]

    class RestoreItem(object):
        def __init__(self, pkg: str):
            self.pkg = pkg
            self.host_apk_list: List[str] = []
            self.users: Union[None, Set[int]] = None
